import matplotlib.pyplot as plt
import os
import csv
from itertools import cycle


class FitPlot:
    def __init__(self, algorithm_names, data, title):
        assert len(algorithm_names) == len(data), "每个算法必须对应一行数据"

        self.algorithm_names = algorithm_names
        self.data = data
        self.title = title
        self.markers = cycle(
            ['o', 's', '*', 'D', '^', 'v', '<', '>', 'p', 'h', '+', 'x', '|', '_'])

    def plot_and_save(self):
        # Create exp_result directory if it doesn't exist
        os.makedirs('exp_result', exist_ok=True)

        plt.figure(figsize=(10, 5))
        plt.title(self.title)
        plt.xlabel('Iterations')
        plt.ylabel('Fitness')

        for algorithm_name, Gbest_curve in zip(self.algorithm_names, self.data):
            marker = next(self.markers)
            iterations = range(1, len(Gbest_curve) + 1)
            plt.plot(iterations, Gbest_curve,
                     label=algorithm_name, marker=marker)

        plt.legend()
        plt.grid()
        # Save the figure before showing it
        plt.savefig(f'exp_result/{self.title}.png')
        plt.show()
        plt.close()

        # Save the ranking to a CSV file
        # Extract the last value from each algorithm's curve
        final_values = [curve[-1] for curve in self.data]
        ranked_algorithms = sorted(
            zip(self.algorithm_names, final_values), key=lambda x: x[1], reverse=True)
        with open(f'exp_result/{self.title}.csv', 'w', newline='') as file:
            writer = csv.writer(file)
            writer.writerow(['Rank', 'Algorithm', 'Final Fitness'])
            for rank, (algorithm_name, fitness) in enumerate(ranked_algorithms, start=1):
                writer.writerow([rank, algorithm_name, fitness])


# # 使用示例
# algorithm_names = ['WOA', 'SSA', 'HHO', 'RIME']
# data = [
#     [945, 8455, 74867, 65646, 545646],  # WOA的数据
#     [95, 475, 44665, 54565, 44565465],  # SSA的数据
#     [77, 888, 7878, 68798, 78958],
#     [789, 8464, 74567, 789896, 4645785],
# ]
# title = 'F10'

# fitness_plot = FitPlot(algorithm_names, data, title)
# fitness_plot.plot_and_save()
